/*
* Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

use std::{collections::{HashMap, VecDeque}, str::FromStr, sync::{Arc, Mutex}, time::Duration};

use dora_message::id::DataId;
use dora_node_api::{arrow::array::{AsArray, UInt8Array}, dora_core::config::NodeId, DoraNode, Event, EventStream};
use dora_tracing::TracingBuilder;
use eyre::Context;

pub struct DoraNodeContext {
    node: DoraNode,
    events: Arc<Mutex<EventStream>>,
    filtered_data: Arc<Mutex<HashMap<String, VecDeque<Vec<u8>>>>>,
}

impl DoraNodeContext {
    pub fn new(node_path: &str) -> eyre::Result<Self> {
        #[cfg(debug_assertions)]
        const FILTER: &str = "debug";
        #[cfg(not(debug_assertions))]
        const FILTER: &str = "warn";

        let node_id = NodeId::from_str(node_path).context("invalid node id")?;
        TracingBuilder::new(node_path)
            .with_stdout(FILTER)
            .build()
            .context("failed to setup tracing")?;
        let (node, events) = DoraNode::init_from_node_id(node_id)
            .context("failed to initialize node")?;
        tracing::info!("node initialized: {}", node_path);
        Ok(Self {
            node,
            events: Arc::new(Mutex::new(events)),
            filtered_data: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    pub fn send(&mut self, path: &str, data: &[u8]) -> eyre::Result<()> {
        let output_id = DataId::from(path.to_string());
        self.node.send_output_raw(output_id, Default::default(), data.len(), |out| {
            out.copy_from_slice(data);
        })
    }

    pub fn recv(&mut self, path: &str, timeout: u32) -> eyre::Result<Vec<u8>> {
        {
            let mut filtered_data = self.filtered_data.lock()
                .map_err(|e| eyre::anyhow!("failed to lock filtered data: {:?}", e))?;
            if let Some(event) = filtered_data.get_mut(path) {
                if let Some(data) = event.pop_front() {
                    return Ok(data);
                }
            }
        }
        
        let data = Self::recv_once(
            self.events.clone(),
            self.filtered_data.clone(),
            path,
            timeout,
        )?;
        if let Some(data) = data {
            return Ok(data);
        }
        Err(eyre::anyhow!("timeout"))
    }

    pub fn on_recv<F>(&mut self, path: &str, callback: F) -> eyre::Result<()>
    where F: Fn(&mut [u8]) + Send + 'static
    {
        let path = path.to_string();
        let events = self.events.clone();
        let filtered_data = self.filtered_data.clone();
        std::thread::spawn(move || {
            loop {
                {
                    let mut filtered_data = match filtered_data.lock() {
                        Ok(filtered_data) => filtered_data,
                        Err(e) => {
                            tracing::error!("failed to lock filtered data: {:?}", e);
                            break;
                        }
                    };
                    if let Some(event) = filtered_data.get_mut(&path) {
                        if let Some(mut data) = event.pop_front() {
                            callback(&mut data);
                            continue;
                        }
                    }
                }

                let data = match Self::recv_once(
                    events.clone(),
                    filtered_data.clone(),
                    &path,
                    u32::MAX,
                ) {
                    Ok(data) => data,
                    Err(error) => {
                        tracing::error!("error: {:?}", error);
                        break;
                    }
                };
                if let Some(mut data) = data {
                    callback(&mut data);
                }
            }
        });
        Ok(())
    }

    fn recv_once(
        events: Arc<Mutex<EventStream>>,
        filtered_data: Arc<Mutex<HashMap<String, VecDeque<Vec<u8>>>>>,
        path: &str,
        timeout: u32,
    ) -> eyre::Result<Option<Vec<u8>>> {
        #[cfg(debug_assertions)]
        tracing::debug!("waiting next event from {} in {}ms", path, timeout);

        let event = events.lock()
            .map_err(|e| eyre::anyhow!("failed to lock events: {:?}", e))?
            .recv_timeout(Duration::from_millis(timeout as u64));
        match event {
            Some(event) => match &event {
                Event::Stop(cause) => Err(eyre::anyhow!("node stopped: {:?}", cause)),
                Event::Reload { .. } => Err(eyre::anyhow!("node reloaded")),
                Event::Input { id, metadata, data } => {
                    match metadata.type_info.data_type {
                        dora_node_api::arrow::datatypes::DataType::UInt8 => {
                            // TODO: optimize this
                            let array: &UInt8Array = data.as_primitive();
                            let data = array.values().to_vec();

                            #[cfg(debug_assertions)]
                            tracing::debug!("received data from {}: {:?}", id.as_str(), data);

                            if id.as_str() == path {
                                Ok(Some(data))
                            } else {
                                let mut filtered_data = filtered_data.lock()
                                    .map_err(|e| eyre::anyhow!("failed to lock filtered data: {:?}", e))?;
                                Self::push_filtered_data(
                                    &mut filtered_data,
                                    id.as_str().to_string(),
                                    data,
                                );
                                Ok(None)
                            }
                        },
                        dora_node_api::arrow::datatypes::DataType::Null => Ok(None),
                        _ => {
                            Err(eyre::anyhow!("unsupported data type: {:?}", metadata.type_info.data_type))
                        }
                    }
                },
                Event::InputClosed { .. } => Err(eyre::anyhow!("input closed")),
                Event::Error(error) => Err(eyre::anyhow!("error: {:?}", error)),
                _ => Err(eyre::anyhow!("unknown event: {:?}", event)),
            },
            None => Ok(None),
        }
    }

    fn push_filtered_data(
        filtered_data: &mut HashMap<String, VecDeque<Vec<u8>>>,
        path: String,
        data: Vec<u8>,
    ) {
        filtered_data.entry(path).or_default().push_back(data);
    }
}
